#!/usr/bin/env python3

import argparse
import matplotlib.pyplot as plt
import numpy as np
import cartopy.crs as ccrs
import xarray as xr
import os
from plot_common import *

parser = argparse.ArgumentParser(description='Plot WRF domains.')
parser.add_argument('--geo-file', '-i', required=True, nargs='+')
args = parser.parse_args()

def plot_domain(ax, geo_file):
	f = xr.open_dataset(geo_file)
	xlon = f.XLONG_M[0,:,:].values
	xlat = f.XLAT_M[0,:,:].values
	dx = f.DX / 1000

	ax.contourf(xlon, xlat, f.HGT_M[0,:,:], cmap='terrain')

	ax.plot(xlon[ 0, :], xlat[ 0, :], transform=ccrs.PlateCarree(), c='blue')
	ax.plot(xlon[-1, :], xlat[-1, :], transform=ccrs.PlateCarree(), c='blue')
	ax.plot(xlon[ :, 0], xlat[ :, 0], transform=ccrs.PlateCarree(), c='blue')
	ax.plot(xlon[ :,-1], xlat[ :,-1], transform=ccrs.PlateCarree(), c='blue')

	ax.text(xlon[-1,0], xlat[-1,0], f'dx={dx}km', fontsize=5, bbox={'facecolor': 'white'})

ax = None
for geo_file in args.geo_file:
	if ax == None:
		f = xr.open_dataset(geo_file)
		min_lon = f.XLONG_M.min()
		max_lon = f.XLONG_M.max()
		min_lat = f.XLAT_M.min()
		max_lat = f.XLAT_M.max()
		dlon = max_lon - min_lon
		dlat = max_lat - min_lat
		ax = get_china_ax(
			min_lon=min_lon-0.1*dlon,
			max_lon=max_lon+0.1*dlon,
			min_lat=min_lat-0.1*dlat,
			max_lat=max_lat+0.1*dlat
		)

	if not os.path.isfile(geo_file):
		print(f'[Error]: File {geo_file} does not exist!')
		exit(1)
	plot_domain(ax, geo_file)

plt.savefig('domains.png', dpi=300)
